package Algorithm.Tree;

import Algorithm.Other.DSU;

import java.util.*;

/**
 * @Author cj
 * @Date 2022/9/8 20:46
 */
public class tarJanQuery {

    class Element<V> {
        public V value;
        public Element(V value) {
            this.value = value;
        }
    }

    class UnionFindSet<V> {
        public HashMap<V, Element<V>> elementMap;
        public HashMap<Element<V>, Element<V>> fatherMap;
        public HashMap<Element<V>, Integer> rankMap;

        public UnionFindSet(List<V> list) {
            elementMap = new HashMap<>();
            fatherMap = new HashMap<>();
            rankMap = new HashMap<>();
            for (V value : list) {
                Element<V> element = new Element<V>(value);
                elementMap.put(value, element);
                fatherMap.put(element, element);
                rankMap.put(element, 1);
            }
        }

        private Element<V> findHead(Element<V> element) {
            Stack<Element<V>> path = new Stack<>();
            while (element != fatherMap.get(element)) {
                path.push(element);
                element = fatherMap.get(element);
            }
            return element;
        }

        public V findHead(V value) {
            return elementMap.containsKey(value) ?
                    findHead(elementMap.get(value).value) : null;
        }

        public boolean isSameSet(V a, V b) {
            if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
                return findHead(elementMap.get(a)) == findHead(elementMap.get(b));
            }
            return false;
        }

        public void union(V a, V b) {
            if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
                Element<V> aF = findHead(elementMap.get(a));
                Element<V> bF = findHead(elementMap.get(b));
                if (aF != bF) {
                    Element<V> big = rankMap.get(aF) > rankMap.get(bF) ? aF : bF;
                    Element<V> small = aF == big ? bF : aF;
                    fatherMap.put(small, big);
                    rankMap.put(big, rankMap.get(big) + rankMap.get(small));
                    rankMap.remove(small);
                }
            }
        }
    }

    class Node {
        public int value;
        public Node left;
        public Node right;

        public Node(int val) {
            this.value = val;
        }
    }

    class Query {
        public Node o1;
        public Node o2;

        public Query(Node o1, Node o2) {
            this.o1 = o1;
            this.o2 = o2;
        }
    }

    public Node[] tarJanQuery1(Node head, Query[] quries) {
        HashMap<Node, LinkedList<Node>> queryMap = new HashMap<>();
        HashMap<Node, LinkedList<Integer>> indexMap = new HashMap<>();
        HashMap<Node, Node> ancestorMap = new HashMap<>();
        UnionFindSet<Node> sets = new UnionFindSet<>(getAllNodes(head));
        Node[] ans = new Node[quries.length];
        setQuriesAndSetEasyAnswers(quries, ans, queryMap, indexMap);
        setAnswers(head, ans, queryMap, indexMap, ancestorMap, sets);
        return ans;
    }

    private List<Node> getAllNodes(Node head) {
        List<Node> res = new ArrayList<>();
        process(head, res);
        return res;
    }

    private void process(Node head, List<Node> res) {
        if (head == null) {
            return;
        }
        res.add(head);
        process(head.left, res);
        process(head.right, res);
    }

    private void setQuriesAndSetEasyAnswers(Query[] quries, Node[] ans, HashMap<Node, LinkedList<Node>> queryMap,
                                              HashMap<Node, LinkedList<Integer>> indexMap) {
        Node o1 = null;
        Node o2 = null;
        for (int i = 0; i < ans.length; i++) {
            o1 = quries[i].o1;
            o2 = quries[i].o2;
            if (o1 == o2 || o1 == null || o2 == null) {
                ans[i] = o1 != null ? o1 : o2;
            } else {
                if (!queryMap.containsKey(o1)) {
                    queryMap.put(o1, new LinkedList<>());
                    indexMap.put(o1, new LinkedList<>());
                }
                if (!queryMap.containsKey(o2)) {
                    queryMap.put(o2, new LinkedList<>());
                    indexMap.put(o2, new LinkedList<>());
                }
                queryMap.get(o1).add(o2);
                indexMap.get(o1).add(i);
            }
        }
    }

    private void setAnswers(Node head, Node[] ans, HashMap<Node, LinkedList<Node>> queryMap,
                            HashMap<Node, LinkedList<Integer>> indexMap, HashMap<Node, Node> ancestorMap,
                            UnionFindSet<Node> sets) {
        if (head == null) {
            return;
        }
        setAnswers(head.left, ans, queryMap, indexMap, ancestorMap,sets);
        sets.union(head.left, head);
        // 设置这个集合的祖先为head
        ancestorMap.put(sets.findHead(head), head);
        // 处理右边
        setAnswers(head.right, ans, queryMap, indexMap, ancestorMap,sets);
        sets.union(head.right, head);
        ancestorMap.put(sets.findHead(head), head);
        LinkedList<Node> nList = queryMap.get(head);
        LinkedList<Integer> iList = indexMap.get(head);
        Node node = null;
        Node nodeFather = null;
        int index = 0;
        while (nList != null && !nList.isEmpty()) {
            node = nList.poll();
            index = iList.poll();
            nodeFather = sets.findHead(node);
            if (ancestorMap.containsKey(nodeFather)) {
                ans[index] = ancestorMap.get(nodeFather);
            }
        }
    }
}
